import sys
import os
import cv2
from PyQt5.QtWidgets import (
    QApplication,
    QMainWindow,
    QLabel,
    QPushButton,
    QFileDialog,
    QListWidget,
    QListWidgetItem,
    QHBoxLayout,
    QVBoxLayout,
    QWidget,
    QSplitter,
    QMessageBox,
    QProgressDialog,
)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt, QSize
from ui.thumbnail_item import ThumbnailItemWidget
from ui.calibration_dialog import CalibrationDialog
from utils.yolov8_detect import detect_beetles, change_model
from utils.model_calibration import calibrate_model
from utils.image_loader import load_images_from_folder, load_single_image


class BeetleCounter(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("天牛计数系统")
        self.resize(1200, 600)

        main_widget = QWidget()
        main_widget.setStyleSheet("background-color: #e0f8e0;")
        main_layout = QHBoxLayout(main_widget)
        self.setCentralWidget(main_widget)

        left_widget = QWidget()
        left_layout = QVBoxLayout(left_widget)

        btn_style = """
            QPushButton {
                font-size: 16px;
                padding: 8px;
                border-radius: 10px;
                background-color: #c0e4c0;
            }
            QPushButton:hover {
                background-color: #a8d4a8;
            }
        """

        self.btn_single = QPushButton("📷 选择单个图像")
        self.btn_single.setStyleSheet(btn_style)
        self.btn_single.clicked.connect(self.load_single_image)

        self.btn_folder = QPushButton("📁 选择文件夹")
        self.btn_folder.setStyleSheet(btn_style)
        self.btn_folder.clicked.connect(self.load_folder)

        self.btn_calibrate = QPushButton("✍️ 手动标定")
        self.btn_calibrate.setStyleSheet(btn_style)
        self.btn_calibrate.clicked.connect(self.calibrate_image)
        self.btn_calibrate.setEnabled(False)

        self.btn_train = QPushButton("🧠 训练模型")
        self.btn_train.setStyleSheet(btn_style)
        self.btn_train.clicked.connect(self.train_model)

        self.btn_select_model = QPushButton("🧩 选择模型")
        self.btn_select_model.setStyleSheet(btn_style)
        self.btn_select_model.clicked.connect(self.select_model)

        btn_group = QVBoxLayout()
        btn_group.addWidget(self.btn_single)
        btn_group.addWidget(self.btn_folder)
        btn_group.addWidget(self.btn_calibrate)
        btn_group.addWidget(self.btn_train)
        btn_group.addWidget(self.btn_select_model)

        left_layout.addLayout(btn_group)

        self.thumbnail_list = QListWidget()
        self.thumbnail_list.setSpacing(5)
        self.thumbnail_list.itemClicked.connect(self.show_image)
        left_layout.addWidget(self.thumbnail_list)

        right_widget = QWidget()
        right_layout = QVBoxLayout(right_widget)

        label_style = "font-size: 16px; padding: 6px;"

        self.original_label = QLabel("原图像")
        self.original_label.setAlignment(Qt.AlignCenter)
        self.original_label.setMinimumSize(320, 240)
        self.original_label.setScaledContents(True)
        self.original_label.setStyleSheet(label_style)

        self.result_label = QLabel("检测后图像")
        self.result_label.setAlignment(Qt.AlignCenter)
        self.result_label.setMinimumSize(320, 240)
        self.result_label.setScaledContents(True)
        self.result_label.setStyleSheet(label_style)

        image_layout = QHBoxLayout()
        image_layout.addWidget(self.original_label)
        image_layout.addWidget(self.result_label)
        right_layout.addLayout(image_layout)

        # ✅ 新增：保存图片按钮
        self.btn_save_image = QPushButton("💾 保存图片")
        self.btn_save_image.setStyleSheet(btn_style)
        self.btn_save_image.clicked.connect(self.save_current_image)
        right_layout.addWidget(self.btn_save_image)

        self.count_label = QLabel("数量：0")
        self.count_label.setAlignment(Qt.AlignCenter)
        self.count_label.setStyleSheet("font-size: 18px; font-weight: bold;")
        right_layout.addWidget(self.count_label)

        self.conf_label = QLabel("平均置信度：0.00")
        self.conf_label.setAlignment(Qt.AlignCenter)
        self.conf_label.setStyleSheet("font-size: 18px; font-weight: bold;")
        right_layout.addWidget(self.conf_label)

        self.conf_list = QListWidget()
        self.conf_list.itemClicked.connect(self.highlight_bbox)
        self.conf_list.setStyleSheet("font-size: 14px;")
        right_layout.addWidget(self.conf_list)

        splitter = QSplitter(Qt.Horizontal)
        splitter.addWidget(left_widget)
        splitter.addWidget(right_widget)
        splitter.setSizes([300, 900])

        main_layout.addWidget(splitter)

        self.image_paths = []
        self.current_image_path = None
        self.current_result_img = None
        self.current_count = 0
        self.bbox_list = []

    def save_current_image(self):
        if self.current_result_img is None or self.current_image_path is None:
            QMessageBox.warning(self, "错误", "无检测结果可保存")
            return

        save_folder = os.path.join("datasets", "wrongImg")
        if not os.path.exists(save_folder):
            os.makedirs(save_folder)

        filename = os.path.basename(self.current_image_path)
        save_path = os.path.join(save_folder, filename)

        try:
            cv2.imwrite(save_path, self.current_result_img)
            QMessageBox.information(self, "成功", f"已保存到 {save_path}")
        except Exception as e:
            QMessageBox.critical(self, "保存失败", f"保存图片出错：{str(e)}")

    def calibrate_image(self):
        if self.current_image_path:
            dialog = CalibrationDialog(self.current_image_path, self)
            if dialog.exec_():
                boxes = dialog.boxes
                if boxes:
                    self.save_calibration_data(boxes)
                else:
                    QMessageBox.warning(self, "警告", "未标记任何天牛位置！")

    def load_single_image(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "选择图像", "", "Images (*.png *.jpg *.jpeg)"
        )
        if path:
            self.image_paths = load_single_image(path)
            self.add_image_to_list(self.image_paths[0])
            self.show_image(self.thumbnail_list.item(self.thumbnail_list.count() - 1))

    def load_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if folder:
            # ✅ 清空原有状态
            self.thumbnail_list.clear()
            self.original_label.clear()
            self.result_label.clear()
            self.count_label.setText("数量：0")
            self.conf_label.setText("平均置信度：0.00")
            self.conf_list.clear()
            self.btn_calibrate.setEnabled(False)
            self.current_image_path = None
            self.current_result_img = None
            self.current_count = 0
            self.bbox_list = []

            # ✅ 加载新文件夹内容
            self.image_paths = load_images_from_folder(folder)
            for image_path in self.image_paths:
                self.add_image_to_list(image_path)

    def add_image_to_list(self, image_path):
        item = QListWidgetItem()
        item.setData(Qt.UserRole, image_path)
        item.setSizeHint(QSize(200, 130))
        widget = ThumbnailItemWidget(image_path)
        self.thumbnail_list.addItem(item)
        self.thumbnail_list.setItemWidget(item, widget)

    def show_image(self, item):
        self.current_image_path = item.data(Qt.UserRole)
        if self.current_image_path:
            try:
                original_img = cv2.imread(self.current_image_path)
                self.display_image(original_img, self.original_label)

                (
                    self.current_result_img,
                    self.current_count,
                    avg_conf,
                    self.bbox_list,
                ) = detect_beetles(self.current_image_path)
                self.display_image(self.current_result_img, self.result_label)

                self.count_label.setText(f"数量：{self.current_count}")
                self.conf_label.setText(f"平均置信度：{avg_conf:.2f}")

                self.conf_list.clear()
                for i, (_, conf) in enumerate(self.bbox_list):
                    self.conf_list.addItem(f"[{i}] 置信度：{conf:.2f}")

                self.btn_calibrate.setEnabled(True)
            except Exception as e:
                self.count_label.setText(f"处理出错: {str(e)}")
                self.conf_label.setText("")
                self.conf_list.clear()
                self.btn_calibrate.setEnabled(False)

    def display_image(self, img, label):
        rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        qt_img = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(qt_img)
        label.setPixmap(pixmap)

    def highlight_bbox(self, item):
        index = self.conf_list.row(item)
        image = cv2.imread(self.current_image_path)

        for i, (xyxy, conf) in enumerate(self.bbox_list):
            x1, y1, x2, y2 = xyxy
            color = (0, 0, 255) if i == index else (0, 255, 0)
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            cv2.putText(
                image,
                f"{conf:.2f}",
                (x1, y1 - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                color,
                2,
            )

        self.display_image(image, self.result_label)

    def save_calibration_data(self, boxes):
        calibration_folder = os.path.join("datasets", "calibration")
        if not os.path.exists(calibration_folder):
            os.makedirs(calibration_folder)

        image_name = os.path.basename(self.current_image_path)
        new_image_path = os.path.join(calibration_folder, image_name)
        cv2.imwrite(new_image_path, cv2.imread(self.current_image_path))

        label_name = os.path.splitext(image_name)[0] + ".txt"
        label_path = os.path.join(calibration_folder, label_name)

        with open(label_path, "w") as f:
            for box in boxes:
                f.write(
                    f"{box[0]} {box[1]:.6f} {box[2]:.6f} {box[3]:.6f} {box[4]:.6f}\n"
                )

        QMessageBox.information(
            self,
            "成功",
            f"标定数据已保存到 {calibration_folder}\n检测到 {len(boxes)} 个天牛",
        )

    def train_model(self):
        progress_dialog = QProgressDialog("正在训练模型，请稍候...", None, 0, 0, self)
        progress_dialog.setWindowTitle("训练中")
        progress_dialog.setCancelButton(None)
        progress_dialog.setMinimumDuration(0)
        progress_dialog.show()
        QApplication.processEvents()
        calibrate_model()
        progress_dialog.close()

    def select_model(self):
        model_path, _ = QFileDialog.getOpenFileName(
            self, "选择模型", "", "Model Files (*.pt)"
        )
        if model_path:
            try:
                change_model(model_path)
                QMessageBox.information(self, "成功", "模型已更新")
            except Exception as e:
                QMessageBox.warning(self, "错误", f"加载模型时出错: {str(e)}")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = BeetleCounter()
    window.show()
    sys.exit(app.exec_())
